
#include "Matrix.hpp"
#include "Poly.hpp"
#include "Util.hpp"
#include "Covariance.hpp"
#include <stdexcept>
#include <fstream>
#include "nr.h"
#include <cstdlib>
#include <unistd.h>

using namespace std;
using namespace Util;

int
main ( int argc, char* argv[], char* env[] )
{
  bool isCovar=false;
  int c;
  while ((c = getopt(argc, argv, "c")) != EOF)
    switch (c)
      { 
      case 'c':
	isCovar = true;
        break;
      case '?':
        cerr << "invalid option" << endl;
        exit(1);
      }

  if ( argc - optind < 1 )
    {
      cerr << "Usage: " << argv[0] << " [-c] <raw data>" << endl;
      return 1;
    }

  try
    {
      //
      // HJM data preparation
      //
      ifstream raw_data(argv[optind]);
      if ( raw_data.fail() )
	{
	  cerr << "open " << argv[optind] << " failed" << endl;
	  exit(1);
	}
      
      // first line is maturities
      string line;
      Vector<double> V_mat;
      getline (raw_data,line);
      istringstream istr(line);
      double tmp;
      while ( istr >> tmp )
	V_mat.push_back(tmp);
      // zero the first maturity
      if ( V_mat[0] < 1.0 )
	V_mat[0] = 0;

      int nmat = V_mat.size();
      cout << nmat << " maturities" << endl;
      //
      // read all the data
      //
      Matrix<double> M_raw;
      Vector<double> Vdat(nmat);
      Vector<double> Vmean(nmat,0.0);
      while ( ! raw_data.eof() )
	{
	  raw_data >> Vdat;
	  M_raw.push_back(Vdat);
	  Vmean += Vdat;
	}
      raw_data.close();
      cout << M_raw.row() << " rows of data read" << endl;

      // mean rates
      Vmean /= M_raw.row();
      Vmean /= 100.0;
      // write initial rates
      ofstream ofs_F0("F0.txt");
      if ( ofs_F0.fail() )
       	{
       	  cerr << "open F0.txt failed" << endl;
       	  exit(1);
       	}
      ofs_F0 << Vmean;
      ofs_F0.close();

      // generate difference matrix
      Matrix<double> Mdiff;
      for ( int i=1; i < M_raw.row(); ++i )
	{
	  Vdat = M_raw[i] - M_raw[i-1];
	  Mdiff.push_back(Vdat);
	}
      //cout << "difference matrix: " << Mdiff.row() << " : " << Mdiff.col() << endl;

      // covariance matrix
      Matrix<double> M_Cov = Mdiff.Tr() * Mdiff;
      M_Cov /= Mdiff.row()-1;
      M_Cov *= 0.0252; // days in a year / BPS
      cout << "covariance matrix: " << M_Cov.row() << " : " << M_Cov.col() << endl;
	  
      // dump the covariance matrix
      ofstream ofs_covar("covar.txt");
      if ( ofs_covar.fail() )
	{
	  cerr << "open \"covar.txt\" failed" << endl;
	  exit(1);
	}
      ofs_covar << M_Cov << endl;
      ofs_covar.close();

      if ( ! M_Cov.isSymmetric() )
	{
	  cerr << "covariance matrix is not symmetric" << endl;
	}

      // remove first value (counter)
      // maturity data are the X values
      int N_X = V_mat.size();
      Vec_DP X(N_X);
      for ( int i=0; i < N_X; ++i )
	X[i] = V_mat[i];

      //cout << cov.row() << " rows, " << cov.col() << " columns" << endl;
      NRMat<double> CovMat = Matrix2NRMat<double>(M_Cov);
      Covariance Covar(CovMat);
      Covar.PCA();
      NRVec<double> EVal = Covar.EigenValues();
      NRVec<int> SV = Covar.SortVec();
      double sum = 0;
      for ( int i = 0; i < EVal.size(); ++i )
	sum += EVal[SV[i]];
      double frc = 0;

      ofstream ofs_eval("eval.txt");
      if ( ofs_eval.fail() )
	{
	  cerr << "open \"eval.txt\" failed" << endl;
	  exit(1);
	}
      ofs_eval << "***** EigenValues *****" << endl;
      //ofs_eval.precision(4);
      //ofs_eval.width(12);
      for ( int i = 0; i < EVal.size(); ++i )
	{
	  frc += EVal[SV[i]];
	  ofs_eval << setw(12) << EVal[SV[i]] << setw(12) << frc/sum << endl;
	}
      ofs_eval.close();

      // Use curve fitting for the top 3 eigenvectors
      //cout << "***** EigenVectors *****" << endl;
      NRMat<double> EVec = Covar.EigenVectors();
      ofstream ofs_vol("vol.txt");
      if ( ofs_vol.fail() )
	{
	  cerr << "open \"volatility.txt\" failed" << endl;
	  exit(1);
	}
      ofstream ofs_evec("evec.txt");
      if ( ofs_evec.fail() )
	{
	  cerr << "open \"evec.txt\" failed" << endl;
	  exit(1);
	}

      ofs_vol << "Maturity";
      for ( int j = 0; j < N_X; ++j )
	ofs_vol << " " << X[j];
      ofs_vol << endl;

      ofs_evec << X[0];
      for ( int j = 1; j < N_X; ++j )
	ofs_evec << " " << X[j];
      ofs_evec << endl;

      // polynomial curve fitting
      // Y data - eigenvectors
      // X data - maturities
      // polynomial of degree 3 - 4 coefficients
      //
      // top 3 eigenvalues
      //
      vector < Poly<double,double> > PVec;
      Matrix<double> EigenVec(3,N_X);
      
      // first eigenvector - constant
      int i = 0;
      //cout << "eigenvector " << i << endl;
      double sqrt_ev = sqrt(EVal[SV[i]]);
      double Vec0Mean = 0;
      Vector<double> Vec0(N_X);
      for ( int j = 0; j < N_X; ++j )
	{
	  double tmp = EVec[j][SV[i]] * sqrt_ev;
	  Vec0[j] = tmp;
	  Vec0Mean += tmp;
	}
      Vec0Mean /= N_X;
      for ( int j = 0; j < N_X; ++j )
	EigenVec[i][j] = Vec0Mean;

      // output volatility data
      ofs_vol << "Vol_" << i;
      for ( int j = 0; j < N_X; ++j )
	ofs_vol << " " << EigenVec[i][j];
      ofs_vol << endl;

      double C0[1];
      C0[0] = Vec0Mean;
      Poly<double,double>P0(C0,0);
      PVec.push_back(P0);

      //cout << "coefficient" << endl;
      //cout << P0;
      ofs_evec << P0;
      ofs_evec << Vec0[0];
      for ( int j = 1; j < N_X; ++j )
	ofs_evec << " " << Vec0[j];
      ofs_evec << endl;

      int N_A = 4;
      for ( i = 1; i < 3; ++i )
	{
	  //cout << "eigenvector " << i << endl;
	  sqrt_ev = sqrt(EVal[SV[i]]);
	  Vec_DP Y(N_X);
	  for ( int j = 0; j < N_X; ++j )
	    Y[j] = EVec[j][SV[i]] * sqrt_ev;

	  // poly of order 3 - 4 coeffs
	  Vec_DP sigma(1.0,N_X);
	  Vec_DP A(N_A);
	  Vec_BOOL IA(true,N_A);
	  Mat_DP A_Covar(N_A,N_A);
	  double chisq;

	  //NR::lfit(X,Y,sigma,A,IA,A_Covar,chisq,&NR::fpoly);
	  Mat_DP U(N_X,N_A),V(N_A,N_A);
	  Vec_DP W(N_A);
	  NR::svdfit(X,Y,sigma,A,U,V,W,chisq,&NR::fpoly);
	  
	  //cout << "coefficients" << endl;
	  //for ( int k=0; k < N_A; ++k )
	  //  cout << " " << A[k];
	  //cout << endl;
	  // check the fit
	  //cout << "ChiSq = " << chisq << endl;

	  PVec.push_back(Poly<double,double>(A.array(),3));
	  Poly<double,double>& PV = PVec[i];
	  //
	  // eigenvector - fitted poly
	  //
	  for ( int j = 0; j < N_X; ++j )
	    EigenVec[i][j] = PV(X[j]);

	  // coefficients
	  for ( int k=0; k < N_A; ++k )
	    ofs_evec << " " << A[k];
	  ofs_evec << endl;
	  // output eigenvectors
	  ofs_evec << Y[0];
	  for ( int j = 1; j < N_X; ++j )
	    ofs_evec << " " << Y[j];
	  ofs_evec << endl;

	  // output volatility data
	  ofs_vol << "Vol_" << i;
	  for ( int j = 0; j < N_X; ++j )
	    ofs_vol << " " << EigenVec[i][j];
	  ofs_vol << endl;
	}
      ofs_evec.close();
      //
      // compute M - integration of the volatility
      //
      // Calculate the polynomial integrals
      vector < Poly<double,double> > PIntegral;
      for ( int i = 0; i < 3; ++i )
	PIntegral.push_back(PVec[i].Integral());

      ofs_vol << "M";
      vector<double> M(N_X,0.0);
      for ( int i=0; i < N_X; ++i )
	{
	  // integral * Vol(tau)
	  for ( int j = 0; j < 3; ++j )
	    M[i] += PIntegral[j].eval(X[i]) * PVec[j].eval(X[i]);
	  ofs_vol << " " << M[i];
	}
      ofs_vol << endl;
      ofs_vol.close();

      //
      // write output data
      //
      ofstream HJM_data("data/HJM.dat",ios::binary);
      if ( HJM_data.fail() )
	{
	  cerr << "open \"data/HJM.dat\" failed" << endl;
	  exit(1);
	}
      HJM_data.write((char*)&N_X,sizeof(int));
      double Vtmp[N_X];
      // write maturities
      for ( int i=0; i < N_X; ++i )
	Vtmp[i] = X[i];
      HJM_data.write((char*)Vtmp,sizeof(double)*N_X); 
      // write drift
      for ( int i=0; i < N_X; ++i )
	Vtmp[i] = M[i];
      HJM_data.write((char*)Vtmp,sizeof(double)*N_X);
      // write eigenvectors
      for ( int i=0; i < 3; ++i )
	{
	  for ( int j=0; j < N_X; ++j )
	    Vtmp[j] = EigenVec[i][j];
	  HJM_data.write((char*)Vtmp,sizeof(double)*N_X);
	}
      HJM_data.close();
    }
  catch ( std::exception& ex )
    {
      cerr << ex.what() << endl;
    }
  exit (0);
}
